package org.zecontech.extraplus.datagen.loot;

import net.minecraft.core.HolderLookup;
import net.minecraft.data.loot.EntityLootSubProvider;
import net.minecraft.world.entity.EntityType;
import net.minecraft.world.flag.FeatureFlags;
import net.minecraft.world.level.ItemLike;
import net.minecraft.world.level.storage.loot.LootPool;
import net.minecraft.world.level.storage.loot.LootTable;
import net.minecraft.world.level.storage.loot.entries.LootItem;
import net.minecraft.world.level.storage.loot.functions.EnchantedCountIncreaseFunction;
import net.minecraft.world.level.storage.loot.functions.SetItemCountFunction;
import net.minecraft.world.level.storage.loot.functions.SmeltItemFunction;
import net.minecraft.world.level.storage.loot.providers.number.ConstantValue;
import net.minecraft.world.level.storage.loot.providers.number.UniformGenerator;
import net.neoforged.neoforge.registries.DeferredHolder;
import org.zecontech.extraplus.init.ModEntities;
import org.zecontech.extraplus.init.ModItems;

import java.util.stream.Stream;

/**
 * 模组实体战利品表提供者
 */
public class ModEntityLoot extends EntityLootSubProvider {

    protected ModEntityLoot(HolderLookup.Provider registries) {
        super(FeatureFlags.DEFAULT_FLAGS, registries);
    }

    /**
     * 生成实体战利品表
     */
    @Override
    public void generate() {
        dropMeat(ModEntities.KELONG.get(), ModItems.KELONG_MEAT);
        dropMeat(ModEntities.NAILONG.get(), ModItems.CHENTAO);
    }

    /**
     * 掉落肉类
     * @param entityType 实体类型
     * @param meat 肉类
     */
    protected void dropMeat(EntityType<?> entityType, ItemLike meat) {
        this.add(entityType, LootTable.lootTable()
                .withPool(LootPool.lootPool()
                        .setRolls(ConstantValue.exactly(1.0F))
                        .add(LootItem.lootTableItem(meat)
                                .apply(SetItemCountFunction.setCount(UniformGenerator.between(1.0F, 3.0F)))
                                .apply(SmeltItemFunction.smelted().when(this.shouldSmeltLoot()))
                                .apply(EnchantedCountIncreaseFunction
                                        .lootingMultiplier(this.registries, UniformGenerator.between(0.0F, 1.0F))))));
    }

    /**
     * 判断是否是已知的实体
     * @return Stream实体列表
     */
    @Override
    protected Stream<EntityType<?>> getKnownEntityTypes() {
        return ModEntities.ENTITIES.getEntries().stream().map(DeferredHolder::get);
    }
}